from typing import List, Tuple

from torch.utils.data import Dataset

from constant import *


def rand_seq(min_len: int = 3, max_len: int = 16) -> List[int]:
    """
    返回一个指定长度范围内的随机整数序列
    """
    # 随机生成序列长度L，在min_len和max_len之间
    L = random.randint(min_len, max_len)
    # 生成L个元素的列表，每个元素是VOCAB_BASE加上0-9之间的随机数
    return [VOCAB_BASE + random.randint(0, 9) for _ in range(L)]


def wrap_tgt(tgt_core: List[int]) -> List[int]:
    """
    接收一个整数列表tgt_core，在其开头添加开始标记BOS_IDX，在末尾添加结束标记EOS_IDX，并将原列表内容反转
    最终返回处理后的新列表
    """
    # tgt_core[::-1] - 将输入列表反转
    # [BOS_IDX] + ... + [EOS_IDX] - 在反转后的列表前后分别添加开始和结束标记
    return [BOS_IDX] + tgt_core[::-1] + [EOS_IDX]


class ReverseDataset(Dataset):
    # 定义了一个名为 ReverseDataset 的数据集类，用于生成用于序列反转任务的样本
    def __init__(self, size: int):
        # 根据给定大小 size，生成多个样本。每个样本包含一个随机整数序列（src）及其反转并添加开始/结束标记的目標序列（tgt），存入 self.samples
        self.samples = []  # 列表
        for _ in range(size):
            src = rand_seq()
            tgt = wrap_tgt(src)
            self.samples.append((src, tgt))

    def __getitem__(self, item: int) -> (List[int], List[int]):
        # 返回指定索引的样本
        return self.samples[item]

    def __len__(self) -> int:
        # 返回样本总数
        return len(self.samples)


"""
zip 函数的双重作用
1.打包功能（常见用法）
将多个可迭代对象按位置组合
例如：zip([1,2], [3,4]) → [(1,3), (2,4)]
2.解包功能（这里使用的用法）
当对已经配对的数据使用时，可以实现"反向"操作
将配对的元素按位置重新分组
"""


def pad_to_max(batch: List[Tuple[List[int], List[int]]]):
    """
    将一个批次中的样本填充到相同长度，并返回填充后的结果
    """
    # *batch对批次进行解包，将每个(sample_src, sample_tgt)元组作为独立参数传递给zip()
    # zip()函数将所有元组的第一个元素（源序列）组合成一个元组，第二个元素（目标序列）组合成另一个元组
    # 最终得到两个独立的序列集合：src_seqs包含所有源序列，tgt_seqs包含所有目标序列
    src_seqs, tgt_seqs = zip(*batch)  # 这是利用了 zip 函数的转置特性，将"行数据"转换为"列数据"
    src_max_len = max(len(seq) for seq in src_seqs)  # 获取源序列中的最大长度
    tgt_max_len = max(len(seq) for seq in tgt_seqs)  # 获取目标序列中的最大长度

    def pad(seqs, max_len):
        """
        填充序列，使其长度与最大长度一致
        """
        out = []
        for seq in seqs:
            # 将序列填充到指定的最大长度
            # seq 是当前需要填充的序列
            # max_len 是批次中所有序列的最大长度
            # PAD_IDX 是填充符号的索引值
            # 在序列末尾添加 (max_len - len(seq)) 个填充符号，使所有序列长度统一
            out.append(seq + [PAD_IDX] * (max_len - len(seq)))
        return torch.tensor(out, dtype=torch.long, device=DEVICE)

    """
    函数返回值默认情况
    Python函数默认返回值是 None
    如果要返回多个值，需要显式使用逗号分隔，这会创建元组
    单个返回值不会自动变成元组
    """
    return pad(src_seqs, src_max_len), pad(tgt_seqs, tgt_max_len)  # 元组


if __name__ == '__main__':
    r_seq = rand_seq()
    # print(r_seq)
    # print(wrap_tgt(r_seq))

    dataset = ReverseDataset(10)
    print(dataset.samples)
    # print(len(dataset))

    # 当使用 for item in dataset 时，Python 会自动调用 dataset.__getitem__(index) 方法
    # 索引从 0 开始递增，直到达到 len 返回的长度
    # 每次迭代返回一个样本元组 (src, tgt)
    # for src, tgt in dataset:
    #     print(src, tgt)

    print(pad_to_max(dataset.samples))
    print(type(pad_to_max(dataset.samples)))
